# 导入必要的库
import os.path

from models import *  # 导入模型相关模块
from utils.datasets import *  # 导入数据集处理模块
from utils.utils import *  # 导入工具函数
import torchvision.transforms as T
from reid.data.transforms import build_transforms  # 行人重识别数据预处理
from reid.modeling import build_model  # 构建行人重识别模型
from reid.config import cfg as reidCfg  # 行人重识别配置
from database import DatabaseManager,VectorDatabase,CacheDatabase


def load_detection_image(self, path):
    """加载检测图像并进行预处理"""
    path = str(Path(path))
    img0 = cv2.imread(path)  # BGR格式

    # 图像预处理
    img, *_ = letterbox(img0, new_shape=self.img_size)
    img = img[:, :, ::-1].transpose(2, 0, 1)  # BGR to RGB, HWC to CHW
    img = np.ascontiguousarray(img, dtype=np.float32)
    img /= 255.0

    # 转换为torch tensor
    img = torch.from_numpy(img).unsqueeze(0).to(self.device)
    return img, img0

class CorridorPedestrianReid:
    def __init__(self, config=None):
        # 模型配置文件路径
        self.cfg = 'cfg/yolov3.cfg'
        # 数据集配置文件所在路径
        self.data = 'data/coco.data'
        # 模型权重文件路径
        self.weights = 'weights/yolov3.weights'
        # 输入分辨率大小
        self.img_size = 416
        # 物体置信度阈值
        self.conf_thres = 0.1
        # NMS阈值
        self.nms_thres = 0.4
        # 行人图片距离阈值，小于这个距离，就认为是该行人
        self.dist_thres = 1.0

        # 初始化设备
        torch.backends.cudnn.benchmark = False
        self.device = torch_utils.select_device(force_cpu=False)

        # 加载类别
        self.classes = load_classes(parse_data_cfg(self.data)['names'])
        self.person_idx = self.classes.index('person')

        # 初始化模型
        self._init_models()

        # 初始化数据库
        self.db = DatabaseManager()
        self.vector_db=VectorDatabase()
        self.cache_db=CacheDatabase(threshold=self.dist_thres)
        # 初始化路径
        self.base_dir = os.path.dirname(os.path.abspath(__file__))

    def _init_models(self):
        """初始化行人检测和重识别模型"""
        with torch.no_grad():
            # 行人检测模型
            self.model = Darknet(self.cfg, self.img_size)
            load_darknet_weights(self.model, self.weights)
            self.model.to(self.device).eval()

            # 行人重识别模型
            self.reid_model = build_model(reidCfg, num_classes=1501)
            self.reid_model.load_param(reidCfg.TEST.WEIGHT)
            self.reid_model.to(self.device).eval()

    def load_detection_image(self,image_bytes):
        """将字节流转换为预处理后的 img 和 img0"""
        # 字节流 -> OpenCV 图像 (BGR 格式)
        nparr = np.frombuffer(image_bytes, np.uint8)
        img0 = cv2.imdecode(nparr, cv2.IMREAD_COLOR)

        # 预处理（与 load_detection_image 一致）
        img, *_ = letterbox(img0, new_shape=self.img_size)  # 调整大小 + 填充
        img = img[:, :, ::-1].transpose(2, 0, 1)  # BGR to RGB, HWC to CHW
        img = np.ascontiguousarray(img, dtype=np.float32)
        img /= 255.0  # 归一化

        # 转换为 PyTorch Tensor（添加 batch 维度）
        img = torch.from_numpy(img).unsqueeze(0)
        return img, img0

    def preprocess_gallery_images(self, gallery_entities):
        """预处理从数据库获取的gallery图像"""
        normalize_transform = T.Normalize(mean=[0.485, 0.456, 0.406],
                                          std=[0.229, 0.224, 0.225])
        transform = T.Compose([
            T.Resize(reidCfg.INPUT.SIZE_TEST),
            T.ToTensor(),
            normalize_transform
        ])

        processed_images = []
        pids = []

        for pid, rel_path in gallery_entities:
            try:
                abs_path = os.path.join(self.base_dir, rel_path)
                if not os.path.exists(abs_path):
                    print(f"Warning: Image file not found for pid {pid}: {abs_path}")
                    continue

                img = Image.open(abs_path).convert('RGB')
                img = transform(img)
                processed_images.append(img)
                pids.append(pid)
            except Exception as e:
                print(f"Error processing image for pid {pid}: {str(e)}")
                continue

        if not processed_images:
            return torch.empty(0), []

        return torch.stack(processed_images, dim=0), pids

    def detect_and_recognize_by_vector_redis_database(self, image_bytes):
        """主函数：检测并识别行人，优先使用Redis查询，其次向量数据库，最后新增"""

        # 加载检测图像
        detected_img, detected_img0 = self.load_detection_image(image_bytes)
        detected_img = detected_img.to(device=self.device)

        with torch.no_grad():
            pred, _ = self.model(detected_img)
            det = non_max_suppression(pred.float(), self.conf_thres, self.nms_thres)[0]

            if det is None or len(det) == 0:
                return detected_img0, [], []

            det[:, :4] = scale_coords(detected_img.shape[2:], det[:, :4], detected_img0.shape).round()

            query_imgs = []
            crop_imgs_to_be_saved = []
            query_locations = []

            for *xyxy, conf, cls_conf, cls in det:
                if int(cls) != self.person_idx:
                    continue

                xmin, ymin, xmax, ymax = map(int, xyxy)
                w, h = xmax - xmin, ymax - ymin
                if w * h > 500:
                    query_locations.append((xmin, ymin, xmax, ymax))
                    crop_img = detected_img0[ymin:ymax, xmin:xmax]
                    crop_imgs_to_be_saved.append(crop_img)

                    # 图像转为模型输入
                    crop_img = Image.fromarray(cv2.cvtColor(crop_img, cv2.COLOR_BGR2RGB))
                    crop_img = build_transforms(reidCfg)(crop_img).unsqueeze(0)
                    query_imgs.append(crop_img)

            if not query_imgs:
                return detected_img0, [], []

            # 提取特征向量
            query_imgs = torch.cat(query_imgs, dim=0).to(self.device)
            query_feats = self.reid_model(query_imgs)
            query_feats = torch.nn.functional.normalize(query_feats, dim=1, p=2)

            # ========== 第一步：Redis 查询 ==========
            redis_matched_ids = self.cache_db.query_persons(query_feats)
            pids = redis_matched_ids.copy()  # 初始化所有结果

            # 统计Redis匹配情况
            redis_matched_count = sum(1 for pid in redis_matched_ids if pid != 0)
            redis_unmatched_count = len(redis_matched_ids) - redis_matched_count
            print(
                f"[Redis查询结果] 总人数: {len(redis_matched_ids)}, 匹配人数: {redis_matched_count}, 未匹配人数: {redis_unmatched_count}")

            # 找出未命中的索引
            to_search_indices = [i for i, pid in enumerate(pids) if pid == 0]
            to_search_feats = query_feats[to_search_indices]


            # ========== 第二步：向量数据库查询 ==========
            if to_search_feats:
                search_results = self.vector_db.search_person(
                    to_search_feats, top_k=1, threshold=self.dist_thres
                )

                # 将匹配结果更新到 pids
                unmatched_indices = []
                unmatched_feats = []
                for local_idx, matched in enumerate(search_results):
                    global_idx = to_search_indices[local_idx]
                    if matched:
                        matched_pid = int(matched[0])
                        pids[global_idx] = matched_pid
                    else:
                        unmatched_indices.append(global_idx)
                        unmatched_feats.append(query_feats[global_idx])

                # ========== 第三步：新行人处理 ==========
                if unmatched_indices:
                    new_imgs = [crop_imgs_to_be_saved[i] for i in unmatched_indices]
                    new_pids = self.db.save_person_image(new_imgs)

                    # 更新向量数据库
                    self.vector_db.add_person(
                        embeddings=torch.stack(unmatched_feats),
                        person_ids=new_pids
                    )

                    # 更新最终PID列表
                    for idx, pid in zip(unmatched_indices, new_pids):
                        pids[idx] = pid

            # ========== 第四步：写入Redis ==========
            self.cache_db.add_persons(pids,query_feats)

            # ========== 第五步：轨迹记录 ==========
            if pids:
                self.db.insert_trajectory_records(pids, camera_id=1)

            return detected_img0, query_locations, pids


if __name__ == '__main__':
    detector = CorridorPedestrianReid()

    # 检测并识别行人
    image_path = "test.jpg"  # 替换为你的图像路径
    original_img, locations, pids = detector.detect_and_recognize_by_vector_redis_database(image_path)

    # 打印结果
    print("检测到的行人位置:", locations)
    print("对应的PID:", pids)

    # 可视化结果
    for (x1, y1, x2, y2), pid in zip(locations, pids):
        cv2.rectangle(original_img, (x1, y1), (x2, y2), (0, 255, 0), 2)
        cv2.putText(original_img, f"PID:{pid}", (x1, y1 - 10),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)

    cv2.imwrite("result.jpg", original_img)
    print("结果已保存为 result.jpg")

